package ml.features;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

import corpus.WordPolarityCorpus;
import corpus.WordPolarity;

import ml.Data;
import parser.ParserAnnotation;

public class WordPolarityExtractor extends FeatureExtractor {
	private WordPolarityCorpus corpus;
	private boolean sequence;

	public WordPolarityExtractor(int offset, WordPolarityCorpus corpus, boolean sequence) {
		super(offset);
		this.corpus = corpus;
		this.sequence = sequence;
	}

	@Override
	public List<Data> extract(ParserAnnotation pa, int sentence) {
		List<String> tokens = pa.getTokens(sentence);
		List<String> pos = pa.getPOS(sentence);

		int nbrPos = 0;
		int nbrNeg = 0;
		int nbrNeu = 0;

		for (int i = 0; i < tokens.size(); i++) {
			WordPolarity wp = corpus.polarity(tokens.get(i), pos.get(i));
			if (wp == null)
				continue;
			if (wp.getEmotion().equals("neg")) {
				if (wp.isStrong())
					nbrNeg++;
				nbrNeg++;
			} else if (wp.getEmotion().equals("pos")) {
				if (wp.isStrong())
					nbrPos++;
				nbrPos++;
			} else {
				if (wp.isStrong())
					nbrNeu++;
				nbrNeu++;
			}
		}
		Data dPos = new Data(nbrPos, 0 + offset);
		Data dNeg = new Data(nbrNeg, 1 + offset);
		Data dNeu = new Data(nbrNeu, 2 + offset);
		List<Data> data = new ArrayList<Data>();
		data.add(dPos);
		data.add(dNeg);
		data.add(dNeu);
		return data;
	}

	public List<Data>[] extractSequence(ParserAnnotation pa, int sentence) {
		List<Data>[] datas = new List[6];
		for (int i = 0; i < 6; i++)
			datas[i] = new ArrayList();
		List<String> sentPos = pa.getPOS(sentence);
		List<String> sentToken = pa.getTokens(sentence);
		WordPolarity[] wps = new WordPolarity[sentToken.size()];
		for (int i = 0; i < wps.length; i++)
			wps[i] = corpus.polarity(sentToken.get(i), sentPos.get(i));
		for (WordPolarity wp : wps) {
			datas[0].add(new Data(wp != null && wp.isPositive(), offset + 0));
			datas[1].add(new Data(wp != null && wp.isNegative(), offset + 1));
			datas[2].add(new Data(wp != null && wp.isNeutral(), offset + 2));
			datas[3].add(new Data(wp != null && wp.isPriorpolarity(), offset + 3));
			datas[4].add(new Data(wp != null && wp.isStrong(), offset + 4));
			datas[5].add(new Data(wp != null && wp.isPriorpolarity(), offset + 5));
		}
		return datas;
	}
	
	@Override
	public int nbrColumns() {
		if (sequence)
			return 6;
		return 3;
	}

	@Override
	public String[] wekaDescription() {
		String pos = "\"DICT POS\" numeric";
		String neg = "\"DICT NEG\" numeric";
		String neu = "\"DICT NEU\" numeric";
		// String ratio = "\"DICT NEG/POS RATIO\" numeric";
		return new String[] { pos, neg, neu };
	}

	@Override
	public void store() throws IOException {
	}
}
